#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
将FrankMocap输出的所有pkl文件合并成一个npz文件
用于Blender导入生成fbx动画
"""

import os
import pickle
import numpy as np
from tqdm import tqdm
import argparse

try:
    from scipy.interpolate import interp1d
    HAS_SCIPY = True
except ImportError:
    HAS_SCIPY = False
    print("Warning: scipy not available, will use simple interpolation")

def load_pkl_frame(pkl_path):
    """加载单个pkl文件并提取关键数据"""
    try:
        with open(pkl_path, 'rb') as f:
            data = pickle.load(f)

        if 'pred_output_list' not in data or len(data['pred_output_list']) == 0:
            return None

        pred_output = data['pred_output_list'][0]

        # 提取所有需要的数据
        frame_data = {}

        # 基础信息
        frame_data['demo_type'] = data.get('demo_type', '')
        frame_data['smpl_type'] = data.get('smpl_type', '')
        
        # 用于标记哪些数据有效
        frame_data['valid_data'] = {}

        # 定义需要插值的数据字段
        interpolatable_fields = [
            ('pred_body_pose', 'poses'),
            ('pred_global_orient', 'global_orient'),
            ('pred_transl', 'trans'),  # 可能不存在，会在后面特殊处理
            ('pred_betas', 'betas'),
            ('pred_left_hand_pose', 'left_hand_pose'),
            ('pred_right_hand_pose', 'right_hand_pose'),
            ('pred_left_hand_betas', 'left_hand_betas'),
            ('pred_right_hand_betas', 'right_hand_betas'),
            ('pred_vertices_smpl', 'vertices'),
            ('pred_camera', 'camera')
        ]

        # 处理可插值的数据字段
        for source_key, target_key in interpolatable_fields:
            if source_key in pred_output and pred_output[source_key] is not None:
                frame_data[target_key] = pred_output[source_key].astype(np.float32)
                frame_data['valid_data'][target_key] = True
            else:
                frame_data['valid_data'][target_key] = False
        
        # 特殊处理：如果没有pred_transl，尝试从pred_camera生成trans
        if 'trans' not in frame_data or not frame_data['valid_data']['trans']:
            if 'pred_camera' in pred_output and pred_output['pred_camera'] is not None:
                camera = pred_output['pred_camera']
                if len(camera) >= 3:
                    # pred_camera格式: [scale, transx, transy]
                    # 生成trans: [transx, transy, 0]
                    trans = np.array([camera[1], camera[2], 0.0], dtype=np.float32)
                    frame_data['trans'] = trans
                    frame_data['valid_data']['trans'] = True
                    print(f"从pred_camera生成trans: {trans}")
                else:
                    # 如果相机参数不足，使用零平移
                    frame_data['trans'] = np.zeros(3, dtype=np.float32)
                    frame_data['valid_data']['trans'] = True
                    print("使用零平移作为默认trans")

        # 面信息 (只需要一次，所有帧都相同)
        if 'faces' in pred_output and pred_output['faces'] is not None:
            frame_data['faces'] = pred_output['faces']

        # 手部方向信息 (也需要插值)
        hand_orient_fields = [
            'left_hand_local_orient_body', 'left_hand_global_orient_body',
            'right_hand_local_orient_body', 'right_hand_global_orient_body',
            'left_hand_local_orient_hand', 'left_hand_global_orient_hand',
            'right_hand_local_orient_hand', 'right_hand_global_orient_hand'
        ]
        
        for key in hand_orient_fields:
            if key in pred_output and pred_output[key] is not None:
                frame_data[key] = pred_output[key].astype(np.float32)
                frame_data['valid_data'][key] = True
            else:
                frame_data['valid_data'][key] = False

        # 关节位置信息 (也需要插值)
        joint_fields = [
            ('pred_body_joints_img', 'body_joints'),
            ('pred_lhand_joints_img', 'left_hand_joints'),
            ('pred_rhand_joints_img', 'right_hand_joints'),
            ('pred_body_joints_3d', 'body_joints_3d'),
            ('pred_lhand_joints_3d', 'left_hand_joints_3d'),
            ('pred_rhand_joints_3d', 'right_hand_joints_3d')
        ]
        
        for source_key, target_key in joint_fields:
            if source_key in pred_output and pred_output[source_key] is not None:
                frame_data[target_key] = pred_output[source_key].astype(np.float32)
                frame_data['valid_data'][target_key] = True
            else:
                frame_data['valid_data'][target_key] = False

        # 现在我们保留所有帧，即使缺少数据，以便后续插值
        return frame_data

    except Exception as e:
        error_msg = str(e)
        if "'NoneType' object has no attribute 'astype'" in error_msg:
            print(f"Error loading {pkl_path}: Data contains None values - this frame likely has incomplete motion capture data")
        else:
            print(f"Error loading {pkl_path}: {e}")
        return None

def interpolate_missing_data(animation_data, all_valid_flags):
    """
    对缺失的数据进行插值处理
    
    Args:
        animation_data: 包含所有动画数据的字典
        all_valid_flags: 每帧每个字段的有效性标记
    """
    num_frames = len(all_valid_flags)
    
    print("Interpolating missing data...")
    interpolated_count = 0
    
    # 获取所有需要插值的字段
    all_fields = set()
    for frame_flags in all_valid_flags:
        if frame_flags:
            all_fields.update(frame_flags.keys())
    
    for field in all_fields:
        if field not in animation_data:
            continue
            
        # 收集有效帧的索引和数据
        valid_indices = []
        valid_data = []
        
        for frame_idx, frame_flags in enumerate(all_valid_flags):
            if frame_flags and frame_flags.get(field, False):
                valid_indices.append(frame_idx)
                valid_data.append(animation_data[field][frame_idx])
        
        if len(valid_indices) < 2:
            # 如果有效数据太少，使用第一个有效值填充所有缺失帧
            if len(valid_indices) == 1:
                fill_value = valid_data[0]
                for frame_idx in range(num_frames):
                    frame_flags = all_valid_flags[frame_idx]
                    if not (frame_flags and frame_flags.get(field, False)):
                        animation_data[field][frame_idx] = fill_value
                        interpolated_count += 1
            continue
            
        # 对多维数据进行插值
        valid_indices = np.array(valid_indices)
        valid_data = np.array(valid_data)
        
        # 找到需要插值的帧
        for frame_idx in range(num_frames):
            frame_flags = all_valid_flags[frame_idx]
            if not (frame_flags and frame_flags.get(field, False)):
                # 进行插值
                if frame_idx < valid_indices[0]:
                    # 前向填充
                    animation_data[field][frame_idx] = valid_data[0]
                elif frame_idx > valid_indices[-1]:
                    # 后向填充
                    animation_data[field][frame_idx] = valid_data[-1]
                else:
                    # 线性插值
                    if HAS_SCIPY:
                        # 使用scipy进行精确插值
                        interp_func = interp1d(valid_indices, valid_data, axis=0, 
                                              kind='linear', bounds_error=False, 
                                              fill_value=(valid_data[0], valid_data[-1]))
                        animation_data[field][frame_idx] = interp_func(frame_idx)
                    else:
                        # 简单线性插值
                        # 找到前后最近的两个有效帧
                        before_idx = np.where(valid_indices < frame_idx)[0]
                        after_idx = np.where(valid_indices > frame_idx)[0]
                        
                        if len(before_idx) > 0 and len(after_idx) > 0:
                            before = before_idx[-1]  # 最近的前一个
                            after = after_idx[0]     # 最近的后一个
                            
                            t = (frame_idx - valid_indices[before]) / (valid_indices[after] - valid_indices[before])
                            interpolated = valid_data[before] * (1 - t) + valid_data[after] * t
                            animation_data[field][frame_idx] = interpolated
                        else:
                            # 如果只有一边有数据，使用最近的有效值
                            if len(before_idx) > 0:
                                animation_data[field][frame_idx] = valid_data[before_idx[-1]]
                            else:
                                animation_data[field][frame_idx] = valid_data[after_idx[0]]
                interpolated_count += 1
    
    if interpolated_count > 0:
        print(f"Interpolated {interpolated_count} missing data points")
    
    return animation_data

def merge_pkl_to_npz(mocap_dir, output_path, frame_range=None):
    """
    将所有pkl文件合并成npz文件

    Args:
        mocap_dir: pkl文件所在目录
        output_path: 输出npz文件路径
        frame_range: 帧范围 (start, end) 或 None (处理所有帧)
    """
    
    # 统计跳过的帧
    skipped_stats = {
        'file_error': 0,
        'missing_data': 0,
        'none_values': 0
    }

    # 获取所有pkl文件
    pkl_files = sorted([f for f in os.listdir(mocap_dir) if f.endswith('.pkl')])

    if not pkl_files:
        print(f"No pkl files found in {mocap_dir}")
        return False

    print(f"Found {len(pkl_files)} pkl files")
    
    # 保存原始文件数量用于统计
    original_file_count = len(pkl_files)

    # 应用帧范围过滤
    if frame_range:
        start_frame, end_frame = frame_range
        if end_frame is not None:
            pkl_files = pkl_files[start_frame:end_frame]
        else:
            pkl_files = pkl_files[start_frame:]
        print(f"Processing frames {start_frame} to {end_frame if end_frame is not None else len(pkl_files)+start_frame-1} ({len(pkl_files)} files)")

    # 初始化数据容器
    animation_data = {}
    faces = None  # 面信息对所有帧都相同
    all_valid_flags = []  # 保存每帧的有效数据标记
    
    # 先处理所有帧，收集数据结构信息
    print("Loading all frames...")
    all_frame_data = []
    for i, pkl_file in enumerate(tqdm(pkl_files, desc="Loading frames")):
        pkl_path = os.path.join(mocap_dir, pkl_file)
        frame_data = load_pkl_frame(pkl_path)

        if frame_data is None:
            print(f"Skipping frame {i}: {pkl_file}")
            skipped_stats['file_error'] += 1
            # 创建空的frame_data以保持帧数
            frame_data = {'valid_data': {}}
            
        all_frame_data.append(frame_data)
        all_valid_flags.append(frame_data.get('valid_data', {}))

    num_frames = len(all_frame_data)
    print(f"Loaded {num_frames} frames, initializing data arrays...")

    # 首先确定所有可能的数据字段（从有效的帧中获取）
    all_data_fields = set()
    demo_type = smpl_type = ""
    
    for frame_data in all_frame_data:
        if 'demo_type' in frame_data:
            demo_type = frame_data['demo_type']
        if 'smpl_type' in frame_data:
            smpl_type = frame_data['smpl_type']
        if 'faces' in frame_data:
            faces = frame_data['faces']
            
        # 收集所有数据字段
        for key, value in frame_data.items():
            if key not in ['demo_type', 'smpl_type', 'faces', 'valid_data'] and isinstance(value, np.ndarray):
                all_data_fields.add(key)
    
    # 设置基础信息
    animation_data['demo_type'] = demo_type
    animation_data['smpl_type'] = smpl_type
    
    # 添加gender字段（Blender SMPL-X插件要求）
    animation_data['gender'] = 'neutral'
    
    # 初始化所有数据数组
    print("Initializing data arrays...")
    for field in all_data_fields:
        # 找到第一个包含此字段的帧来确定数据形状
        for frame_data in all_frame_data:
            if field in frame_data:
                value = frame_data[field]
                shape = (num_frames,) + value.shape
                animation_data[field] = np.zeros(shape, dtype=value.dtype)
                break
    
    # 填充数据
    print("Filling data arrays...")
    for i, frame_data in enumerate(tqdm(all_frame_data, desc="Filling arrays")):
        for field in all_data_fields:
            if field in frame_data and field in animation_data:
                animation_data[field][i] = frame_data[field]

    # 添加面信息
    if faces is not None:
        animation_data['faces'] = faces

    # 进行插值处理以填充缺失数据
    animation_data = interpolate_missing_data(animation_data, all_valid_flags)
    
    # 检查数据质量
    total_missing = sum(skipped_stats.values())
    if total_missing > 0:
        print(f"Warning: {total_missing} frames had data issues, but were interpolated")
    
    # 检查是否帧数太少
    if num_frames < 10:
        print(f"Warning: Very few frames ({num_frames}) - animation may be too short")

    # 添加元数据
    animation_data['num_frames'] = num_frames
    animation_data['fps'] = 30.0  # FrankMocap默认30fps
    animation_data['total_duration'] = num_frames / 30.0

    # 创建输出目录
    output_dir = os.path.dirname(output_path)
    if output_dir:  # 只有当输出路径包含目录时才创建
        os.makedirs(output_dir, exist_ok=True)

    # 保存为npz文件
    print(f"Saving animation data to {output_path}")
    try:
        np.savez_compressed(output_path, **animation_data)
    except Exception as e:
        print(f"Error saving npz file: {e}")
        return False

    # 打印摘要信息
    print("\n=== Animation Data Summary ===")
    print(f"Total frames: {num_frames}")
    print(f"Duration: {num_frames/30.0:.2f} seconds")
    
    # 显示数据质量统计
    if skipped_stats['file_error'] > 0:
        print(f"Frames with missing data: {skipped_stats['file_error']} ({skipped_stats['file_error']*100/num_frames:.1f}%)")
        print("Missing data was interpolated from neighboring frames")
    
    print(f"Data arrays saved:")

    for key, value in animation_data.items():
        if isinstance(value, np.ndarray):
            print(f"  {key}: {value.shape} ({value.dtype})")
        else:
            print(f"  {key}: {type(value).__name__} = {value}")

    # 获取文件大小（仅在文件存在时）
    try:
        file_size = os.path.getsize(output_path) / (1024 * 1024)  # MB
        print(f"\nOutput file size: {file_size:.2f} MB")
    except OSError:
        print("\nCould not determine file size")

    return True

def main():
    parser = argparse.ArgumentParser(description='Merge FrankMocap pkl files to npz format')
    parser.add_argument('--input', '-i',
                       default='./mocap_output_full/mocap',
                       help='Input directory containing pkl files')
    parser.add_argument('--output', '-o',
                       default='./output_npz/smplx_animation.npz',
                       help='Output npz file path')
    parser.add_argument('--start-frame', type=int, default=None,
                       help='Start frame index (0-based)')
    parser.add_argument('--end-frame', type=int, default=None,
                       help='End frame index (exclusive)')

    args = parser.parse_args()

    # 检查输入目录
    if not os.path.exists(args.input):
        print(f"Error: Input directory {args.input} does not exist")
        return

    # 设置帧范围
    frame_range = None
    if args.start_frame is not None or args.end_frame is not None:
        start = args.start_frame if args.start_frame is not None else 0
        end = args.end_frame  # None表示到最后一帧
        frame_range = (start, end)
        print(f"Frame range specified: {start} to {end if end is not None else 'end'}")

    # 执行合并
    print("=== FrankMocap PKL to NPZ Converter ===")
    print(f"Input directory: {args.input}")
    print(f"Output file: {args.output}")

    success = merge_pkl_to_npz(args.input, args.output, frame_range)

    if success:
        print("\n✅ Conversion completed successfully!")
        print(f"Animation data saved to: {args.output}")
        print("\nKey improvements in this version:")
        print("- All frames preserved (no skipping)")
        print("- Missing data interpolated for smooth animation")
        print("- Maintains original timeline and duration")
        print("\nYou can now use this npz file in Blender for animation generation.")
    else:
        print("\n❌ Conversion failed!")

if __name__ == "__main__":
    main()